fix(neighbors): harden selective cluster-tile state - #141
Open
zubatyuk wants to merge 6 commits into
Open
Conversation
Dispatch the existing wrapped PBC kernels after position wrapping for normal and selective rebuilds. Add parity coverage against single-cutoff results for both dtypes and wrap modes. Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
Validate fixed-capacity COO state before launches, preserve JAX empty selective state, and support fullgraph single-system COO reuse. Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile SummaryThis PR hardens selective cluster-tile neighbor-list state handling across PyTorch and JAX.
Important Files Changed
Reviews (4): Last reviewed commit: "Merge branch 'main' into fix/cluster-til..." | Re-trigger Greptile |
Prevent malformed segmented COO metadata from writing past physical output capacity. Document the compiled cluster-tile requirement for eagerly validated, fully periodic PBC. Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
laserkelvin
reviewed
Jul 30, 2026
laserkelvin
left a comment
laserkelvin
left a comment
Collaborator
There was a problem hiding this comment.
Several correctness and state-lifecycle issues remain; details are inline. F1 already had an existing inline thread, so I did not duplicate it.
Add eager COO bootstrap, validate fixed-state buffers before launch, and reject compiled unified dispatch that cannot validate tensor-valued PBC. Cover compiled reuse, CUDA graph capture, and wrapped-PBC selective behavior. Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
Comment on lines
+239
to
+247
| if not torch.compiler.is_compiling(): | ||
| _validate_segmented_coo_values( | ||
| neighbor_list=neighbor_list, | ||
| pair_offsets=pair_offsets, | ||
| pair_counts=pair_counts, | ||
| tile_offsets=tile_offsets, | ||
| tile_counts=tile_counts, | ||
| tile_row_group=tile_row_group, | ||
| ) |
There was a problem hiding this comment.
When a compiled selective COO call receives a reversed or undersized segment, this branch skips the value validation and the kernel suppresses writes for the invalid interval while still incrementing pair_counts. The returned count then identifies unwritten entries as active topology, leaving the previously reported malformed-metadata failure reachable.
Knowledge Base Used: Neighbor list computation (Warp core)
laserkelvin
approved these changes
Jul 31, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
ALCHEMI Toolkit-Ops Pull Request
Description
Harden selective cluster-tile neighbor-list state handling across PyTorch and JAX, including fixed-capacity COO validation, empty-system semantics, and compiled single-system reuse.
Type of Change
Related Issues
None.
Changes Made
Testing
make pytest)make lint)Checklist
Additional Notes
None.
Tip
This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.